'''
@Author: JintuZheng 郑晋图
@Code: Run demo for image
'''
import torch
from torchvision import transforms
from skimage import io
from model import R34
import numpy as np

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu') # Device set

#img = io.imread("Dataset/mnist_test/3/mnist_test_32.png")
img = io.imread('Dataset/mnist_test_6.png')
tfunc = transforms.ToTensor()
img = tfunc(img)
img = torch.unsqueeze(img, dim=0) #
img = img.to(device)
model = R34(10)
model = model.to(device)

model.load_state_dict(torch.load('weights/resnet34_mnist_5.pth'))
model.eval()

pred = model(img)
pred = pred.detach().to('cpu').numpy() 
idx = np.argmax(pred, axis=1) # class-idxs #result
print('The number is {}'.format(idx))
